import hashlib
import os
from sshtunnel import SSHTunnelForwarder
import configparser


class SSHTunnelManager:
    def __init__(self, config_file='config.ini', callback=None):
        self.callback = callback
        self.server_info, self.targets = self.parse_config(config_file)
        self.tunnels = [self.create_tunnel(target) for target in self.targets]

    def print_msg(self, msg):
        if self.callback:
            self.callback(msg)
        else:
            print(msg)

    @staticmethod
    def calculate_md5(file_path):
        with open(file_path, 'rb') as f:
            return hashlib.md5(f.read()).hexdigest()

    def parse_config(self, config_file):
        msg = f"""Warning in [parse_config]: 
您还未编辑配置文件, 请编辑通目录下的{config_file}后再开始隧道服务
配置文件格式如下: 
[ServerA]: 填写ssh服务器的地址
[Target]:  填写本地IP以及本地端口,  还有需要访问的目标机器及端口
[Target] 可以添加多个"""
        config = configparser.ConfigParser()
        if not os.path.exists(config_file):
            # 添加ServerA部分
            config['Server'] = {
                'ip': '192.168.123.228',
                'port': '22',
                'username': 'root',
                'password': '******',
                'use_ssh_key': 'False',
                'ssh_key_file': 'D:\\root.pem',
                'ssh_private_key_password': '',
                'remarks': '上述配置信息中, 如果使用秘钥登录,则将use_ssh_key设置为True, 并传入秘钥的路径及秘钥密码(如果有)'
            }
            # 添加Target1部分
            config['Target1'] = {
                'remarks': '172.168.123.228的数据库服务(备注信息,必填)',
                'target_ip': '172.168.123.228',
                'target_port': '3306',
                'local_listen_ip': '127.0.0.1',
                'local_listen_port': '3306'
            }
            with open('config.ini', 'w', encoding='UTF-8') as configfile:
                config.write(configfile)
            raise Exception(msg)
        else:
            if self.calculate_md5(config_file) == "2bf7f9d0fc8664e871f4bff66f44d4b8":
                raise Exception(msg)
        config.read(config_file, encoding='UTF-8')
        server_info = {
            "ssh_address_or_host": config.get('Server', 'ip'),
            "ssh_port": config.getint('Server', 'port'),
            "ssh_username": config.get('Server', 'username')
        }

        use_ssh_key = config.getboolean('Server', 'use_ssh_key')
        if use_ssh_key:
            ssh_key_file = config.get('Server', 'ssh_key_file')
            ssh_private_key_password = config.get('Server', 'ssh_private_key_password')
            server_info["ssh_pkey"] = ssh_key_file
            if ssh_private_key_password:
                server_info["ssh_private_key_password"] = ssh_private_key_password
        else:
            server_info["ssh_password"] = config.get('Server', 'password')

        targets = []
        for section in config.sections():
            if section.startswith('Target'):
                target = {
                    "remarks": config.get(section, 'remarks'),
                    "remote_bind_address": (config.get(section, 'target_ip'), config.getint(section, 'target_port')),
                    "local_bind_address": (
                        config.get(section, 'local_listen_ip'), config.getint(section, 'local_listen_port')),
                }
                targets.append(target)
        return server_info, targets

    def create_tunnel(self, target):
        if "ssh_password" in self.server_info:
            tunnel = SSHTunnelForwarder(
                (self.server_info["ssh_address_or_host"], self.server_info["ssh_port"]),
                ssh_username=self.server_info["ssh_username"],
                ssh_password=self.server_info["ssh_password"],
                remote_bind_address=target["remote_bind_address"],
                local_bind_address=target["local_bind_address"]
            )
        elif "ssh_pkey" in self.server_info and "ssh_private_key_password" not in self.server_info:
            tunnel = SSHTunnelForwarder(
                (self.server_info["ssh_address_or_host"], self.server_info["ssh_port"]),
                ssh_username=self.server_info["ssh_username"],
                remote_bind_address=target["remote_bind_address"],
                local_bind_address=target["local_bind_address"],
                ssh_pkey=self.server_info["ssh_pkey"]
            )
        else:
            tunnel = SSHTunnelForwarder(
                (self.server_info["ssh_address_or_host"], self.server_info["ssh_port"]),
                ssh_username=self.server_info["ssh_username"],
                remote_bind_address=target["remote_bind_address"],
                local_bind_address=target["local_bind_address"],
                ssh_pkey=self.server_info["ssh_pkey"],
                ssh_private_key_password=self.server_info["ssh_private_key_password"]
            )

        return tunnel

    def start_tunnel(self, tunnel, target):
        tunnel.start()
        msg = f"启动隧道 {target['remarks']}: {target['local_bind_address'][0]}:{target['local_bind_address'][1]} -> {target['remote_bind_address'][0]}:{target['remote_bind_address'][1]} "
        self.print_msg(msg)

    def start_all_tunnels(self):
        for tunnel, target in zip(self.tunnels, self.targets):
            self.start_tunnel(tunnel, target)

    def stop_all_tunnels(self):
        for tunnel in self.tunnels:
            tunnel.stop()


if __name__ == "__main__":
    tunnel_manager = SSHTunnelManager()
    tunnel_manager.start_all_tunnels()

    # 在这里添加代码以便在需要时调用tunnel_manager.stop_all_tunnels()来停止所有SSH隧道
